import numpy as np
import my_assert as my
import random as rand
import time
import math as m

def unittest_ufunc_reduce():
	print "unittest_ufunc_reduce() - starting"
	#import numpy as np
	rand_low = 150
	rand_high = 250
	#1d array
	cuda_A = np.zeros((rand.randrange(rand_low,rand_high)), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	numpy_A = my.copy_1d(cuda_A)
	cuda_res = np.add.reduce(cuda_A)
	numpy_res = np.add.reduce(numpy_A)
	c = my.assert_scalar(cuda_res, numpy_res)
	if c > 0:
		print "1d add.reduce failed"
		print "cuda_A:",cuda_A
		return
	print '2D'
	#2d array
	cuda_A = np.zeros((rand.randrange(rand_low,rand_high), rand.randrange(rand_low,rand_high)), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	numpy_A = my.copy_2d(cuda_A)
	cuda_res = np.add.reduce(cuda_A)
	numpy_res = np.add.reduce(numpy_A)
	c = my.assert_ary_1d(cuda_res, numpy_res)
	if c > 0:
		print "2d add.reduce failed"
		print "shape(cuda_A):", np.shape(cuda_A)
		return
	print '2D axis=1'
	cuda_res = np.add.reduce(cuda_A, axis=1)
	numpy_res = np.add.reduce(numpy_A, axis=1)
	c = my.assert_ary_1d(cuda_res, numpy_res)
	if c > 0:
		print "2d add.reduce failed for axis=1"
		print "shape(cuda_A):", np.shape(cuda_A)
		return
	print '3D'
	#3d array	
	cuda_AA = np.zeros((rand.randrange(rand_low,rand_high), rand.randrange(rand_low,rand_high), rand.randrange(rand_low,rand_high)), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_AA, cuda_AA)	
	numpy_A = my.copy_3d(cuda_AA)
	cuda_res = np.add.reduce(cuda_AA)	
	numpy_res = np.add.reduce(numpy_A)
	c = my.assert_ary_2d(cuda_res, numpy_res)
	if c > 0:
		print "3d add.reduce failed"
		print "shape(cuda_AA):", np.shape(cuda_AA)
		return
	cuda_res = np.add.reduce(cuda_AA, axis=2)
	numpy_res = np.add.reduce(numpy_A, axis=2)
	c = my.assert_ary_2d(cuda_res, numpy_res)
	if c > 0:
		print "3d add.reduce failed for axis=2"
		print "shape(cuda_AA):", np.shape(cuda_AA)
		return

	print "unittest_ufunc_reduce() - completed"

def unittest_ufunc():
	print "unittest_ufunc() - starting"
	#Special cases
	cuda_A = np.array([1,2,3,4], cuda=1)
	numpy_A = np.array([1,2,3,4], cuda=0)
	cuda_res = np.add(cuda_A,1)
	numpy_res = np.add(numpy_A,1)
	c=my.assert_ary_1d(cuda_res, numpy_res)
	if c>0:
		print 'unittest_ufunc - failed vector+scalar'
	cuda_res = np.add(1,cuda_A)
	numpy_res = np.add(1,numpy_A)

	c=my.assert_ary_1d(cuda_res, numpy_res)	

	if c>0:
		print 'unittest_ufunc - failed scalar+vector'
	cuda_res = np.square(cuda_A)
	numpy_res = np.square(numpy_A)
	c=my.assert_ary_1d(cuda_res, numpy_res)	
	if c>0:
		print 'unittest_ufunc - failed a^2'
	print 'unittest_ufunc() - special cases done'
	#General cases
	a=np.zeros((16,4))
	x=0
	for i in [1,6]:
		for j in [1,4]:
			for k in [1,5]:
				for l in [1,7]:
					a[x]=[i,j,k,l]
					x+=1
	b=a
	for i in range(0,16):#a
		print "i:",i
		for j in range(0,16):#b
			cuda_A = np.zeros((a[i][0], a[i][1], a[i][2], a[i][3]), cuda=1)
			cuda_B = np.zeros((b[j][0], b[j][1], b[j][2], b[j][3]), cuda=1)
			#print 'shape cuda_A:',np.shape(cuda_A), 'cuda_B:',np.shape(cuda_B)
			np.ufunc_random(cuda_A, cuda_A)
			np.ufunc_random(cuda_B, cuda_B)
			numpy_A = my.copy_4d(cuda_A)
			numpy_B = my.copy_4d(cuda_B)
			#ufunc add
			cuda_res = np.add(cuda_A, cuda_B)
			numpy_res = np.add(numpy_A, numpy_B)
			c_add = my.assert_ary_4d(cuda_res, numpy_res)
			#ufunc subtract
			cuda_res = np.subtract(cuda_A, cuda_B)
			numpy_res = np.subtract(numpy_A, numpy_B)
			c_sub = my.assert_ary_4d(cuda_res, numpy_res)
			if c_add > 0 or c_sub > 0:
				c = my.assert_ary_4d(cuda_A, numpy_A)
				if c >0:
					print "cuda_A != numpy_A"
				c = my.assert_ary_4d(cuda_B, numpy_B)
				if c >0:
					print "cuda_B != numpy_B"
				print "np.shape(cuda_A):", np.shape(cuda_A),"cuda_A: ", cuda_A 
				print "np.shape(cuda_B):", np.shape(cuda_B),"cuda_B: ", cuda_B 
				print "cuda_res: ", cuda_res
				print "numpy_res: ", numpy_res
				return
	print "unittest_ufunc() - completed"

def unittest_dot():	
	print 'unittest_dot()- starting'
	#Asynchronous matrix dot matrix
	cuda_A = np.zeros((12, 4), dtype=np.double, cuda=1)
	cuda_B = np.zeros((4, 2), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	np.ufunc_random(cuda_B, cuda_B)
	numpy_A = my.copy_2d(cuda_A)
	numpy_B = my.copy_2d(cuda_B)
	cuda_res=np.dot(cuda_A, cuda_B)
	numpy_res=np.dot(numpy_A, numpy_B)
	c=my.assert_ary_2d(cuda_res, numpy_res)
	if c>0:
		if 0<my.assert_ary_2d(cuda_A, numpy_A):
			print 'cuda_A != numpy_A'
		if 0<my.assert_ary_2d(cuda_B, numpy_B):
			print 'cuda_B != numpy_B'
		print 'unittest_dot failed - Asynchronous!'
		return
	#Asynchronous matrix dot matrix
	cuda_A = np.zeros((8, 11), dtype=np.double, cuda=1)
	cuda_B = np.zeros((11, 14), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	np.ufunc_random(cuda_B, cuda_B)
	numpy_A = my.copy_2d(cuda_A)
	numpy_B = my.copy_2d(cuda_B)
	cuda_res=np.dot(cuda_A, cuda_B)
	numpy_res=np.dot(numpy_A, numpy_B)
	c=my.assert_ary_2d(cuda_res, numpy_res)
	if c>0:
		if 0<my.assert_ary_2d(cuda_A, numpy_A):
			print 'cuda_A != numpy_A'
		if 0<my.assert_ary_2d(cuda_B, numpy_B):
			print 'cuda_B != numpy_B'
		print 'unittest_dot failed - Asynchronous!'
		return	
	#Synchronous matrix dot matrix
	cuda_A = np.zeros((32, 32), dtype=np.double, cuda=1)
	cuda_B = np.zeros((32, 32), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	np.ufunc_random(cuda_B, cuda_B)
	numpy_A = my.copy_2d(cuda_A)
	numpy_B = my.copy_2d(cuda_B)
	cuda_res=np.dot(cuda_A, cuda_B)
	numpy_res=np.dot(numpy_A, numpy_B)
	c=my.assert_ary_2d(cuda_res, numpy_res)
	if c>0:
		if 0<my.assert_ary_2d(cuda_A, numpy_A):
			print 'cuda_A != numpy_A'
		if 0<my.assert_ary_2d(cuda_B, numpy_B):
			print 'cuda_B != numpy_B'
		print 'unittest_dot failed - Synchronous!'
		return
	cuda_A = np.zeros((1000, 1000), dtype=np.double, cuda=1)
	cuda_B = np.zeros((1000, 1000), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	np.ufunc_random(cuda_B, cuda_B)
	#print 'a'
	numpy_A = my.copy_2d(cuda_A)
	numpy_B = my.copy_2d(cuda_B)
	#print 'b'
	t0 = time.time()
	cuda_res=np.dot(cuda_A, cuda_B)
	t1 = time.time()
	#print 'cuda:',t1-t0
	numpy_res=np.dot(numpy_A, numpy_B)
	t2 = time.time()
	print 'cuda:',t1-t0,'numpy:',t2-t1
	c=my.assert_ary_2d(cuda_res, numpy_res)
	if c>0:
		if 0<my.assert_ary_2d(cuda_A, numpy_A):
			print 'cuda_A != numpy_A'
		if 0<my.assert_ary_2d(cuda_B, numpy_B):
			print 'cuda_B != numpy_B'
		print 'unittest_dot failed - Synchronous!'
		return

	print 'unittest_dot()- completed!'

def unittest_2d_fft(x=4, y=4):
	#print 'unittest_2d_fft - starting'
	c=0
	max_e=0
	cuda_A  = np.zeros((x,y), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	cuda_B = np.cufftR2C(cuda_A)
	numpy_B = my.copy_2d(cuda_B)
	#Testing complex to complex
	cuda_C = np.cufftC2C(cuda_B)
	numpy_C = np.fft.fftn(numpy_B)
	for i in np.arange(x):
		for j in np.arange(y):
			e1 = m.fabs(cuda_C[i][j].real - numpy_C[i][j].real)
			e2 = m.fabs(cuda_C[i][j].imag - numpy_C[i][j].imag)
			#Testing with an error tolerance of +- 0.001 %
			if e1>m.fabs((numpy_C[i][j].real)/100000) or e2>m.fabs((numpy_C[i][j].imag)/100000):
				if e1 > max_e:
					max_e = e1
				if e2 > max_e:
					max_e = e2
				if c == 0:
					print 'forward: cuda_C[',i,'][',j,']:', cuda_C[i][j], 'numpy_C[',i,'][',j,']:', numpy_C[i][j]
				c+=1
	#Testing inverse complex to complex
	cuda_D = np.cufftiC2C(cuda_C)
	numpy_D = np.fft.ifftn(numpy_C)*(x*y)
	for i in np.arange(x):
		for j in np.arange(y):
			e1 = m.fabs(cuda_D[i][j].real - numpy_D[i][j].real)
			e2 = m.fabs(cuda_D[i][j].imag - numpy_D[i][j].imag)
			#Testing with an error tolerance of +- 0.001 %
			if e1>m.fabs((numpy_D[i][j].real)/10000) or e2>0.0001:
				if e1 > max_e:
					max_e = e1
				if e2 > max_e:
					max_e = e2
				if c == 0:
					print 'inverse: cuda_D[',i,'][',j,']:', cuda_D[i][j], 'numpy_D[',i,'][',j,']:', numpy_D[i][j]
				c+=1
	#Testing fft(ifft(cuda_A))/length(cuda_A)==cuda_A
	cuda_C = np.cufftC2R(cuda_B)/(x*y)
	c += my.assert_ary_2d(cuda_A, cuda_C)
	if max_e > 0:
		print max_e
	#print 'unittest_2d_fft - completed'
	return c

def unittest_1d_fft(x=333):
	#np.random.uniform(0, 2*np.pi, (20,))*1j #Random complex doubles
	#print 'unittest_1d_fft() - starting'
	cuda_A = np.zeros((x), dtype=np.double, cuda=1)
	np.ufunc_random(cuda_A, cuda_A)
	numpy_A = my.copy_1d(cuda_A)
	#real -> complex
	cuda_res = np.cufftR2C(cuda_A)
	numpy_res = np.fft.fft(numpy_A)
	for i in np.arange(x):
		if m.fabs(cuda_res[i].real-numpy_res[i].real)>0.000001 or m.fabs(cuda_res[i].imag-numpy_res[i].imag)>0.000001:
			print 'error i:',i,'x:',x
			print 'cuda_res[i]', cuda_res[i]
			print 'numpy_res[i]', numpy_res[i]
			return 0
		i+=1
	#complex -> complex
	ct = np.cufftC2C(cuda_res)
	nt = np.fft.fft(numpy_res)
	for i in np.arange(x):
		if m.fabs(ct[i].real-nt[i].real)>0.000001 or m.fabs(ct[i].imag-nt[i].imag)>0.000001:
			print 'error i:',i,'x:',x
			print 'ct[i]', ct[i]
			print 'nt[i]', nt[i]
			return 0
		i+=1
	#complex -> real
	ct = (np.cufftC2R(cuda_res))/x
	nt = np.fft.ifft(numpy_res)
	for i in np.arange(x):
		if m.fabs(ct[i]-nt[i].real)>0.000001:
			print 'error i:',i,'x:',x
			return 0
		i+=1
	#print 'unittest_1d_fft() - completed!'
	return 1

def time_test():
	d0 = 10
	d1 = 50
	d2 = 100
	d3 = 70
	ca = np.zeros((d0, d1, d2, d3), dtype=np.double, cuda=1)
	cb = np.zeros((d0, d1, d2, d3), dtype=np.double, cuda=1)
	na = np.zeros((d0, d1, d2, d3), dtype=np.double)
	nb = np.zeros((d0, d1, d2, d3), dtype=np.double)
	np.ufunc_random(ca, ca)
	np.ufunc_random(cb, cb)
	np.ufunc_random(na, na)
	np.ufunc_random(nb, nb)
	t0 = time.time()
	np.add(ca,cb)
	t1 = time.time()
	print 'cuda time: ',t1-t0
	t0 = time.time()
	np.add(ca,cb)
	t1 = time.time()
	print 'cuda time: ',t1-t0
	t0 = time.time()
	np.add(na,nb)
	t1 = time.time()
	print 'numpy time:',t1-t0
	return

unittest_dot()
unittest_ufunc_reduce()
unittest_ufunc()
unittest_2d_fft(61,63)
unittest_1d_fft(33)

print 'unittest_fft() - starting'
for i in np.arange(3,2000,153):
	e1=unittest_1d_fft(i)
	if not(e1):
		print 'i:',i,'errors',e1
		break
	else:
		for j in np.arange(3,300, 7):
			e2=unittest_2d_fft(i,j)
			if e2:
				print 'i:',i,'j:',j,'errors:',e2
				break			
print 'unittest_fft() - completed'

#time_test()
